#-------------------------------------------------------------------------------
# Name:        blasttofasta
# Purpose:     Converts blast xml output to fasta format
#
# Author:      Bc. Ondrej Pelikan
#
# Created:     05.04.2013
#-------------------------------------------------------------------------------

import logging
from xml.etree import ElementTree



def convertXMLtoFASTA(inputStr, logExcluded=True):
    """
    Parameter inputStr is string from xml tree is reading.
    Returns None if error occured. Or retuns string contains FASTA sequences, but only
        sequences that contains exactly standard 20 amino acid letters. Sequence
        containing some nonstandard letter are excluded.
    Raises ElementTree.ParseError when errors occure in xml tree.
    """
    from cStringIO import StringIO
    import re
    import mappanalysis as mapp
    


    #convert xml to fasta and write to the output
    output = StringIO()
    checkReg = re.compile("^[" + "".join(mapp.aalist) + "-" + "]+$", re.IGNORECASE)
    excluded = 0
##    try:
    root = ElementTree.fromstring(inputStr)
    for hit in root.getiterator('Hit'):
        hit_seq = hit.find('Hit_hsps/Hsp/Hsp_hseq')
        hit_seq.text = re.compile("[ -]").sub("", hit_seq.text)
        if (checkReg.match(hit_seq.text)):
            hit_id = hit.find('Hit_id')
            idstr = "gi_" + hit_id.text.split("|")[1]
            output.write('>' + idstr + '\n')
            output.write(hit_seq.text + '\n')
        else:
            excluded += 1
    if logExcluded:
        message = "XMLtoFASTA: " + str(excluded) + " sequences was excluded from BLAST output.\n"
        if __name__ == "__main__":
            sys.stdout.write(message)
        else:
            logger = logging.getLogger(__name__)
            logger.warning(message)
    return output.getvalue()

##    except ElementTree.ParseError: # ParseError is not defined in python 2.5
##        raise


"""
MAIN FUNCTION
"""
if __name__ == "__main__":
    import sys
    import argparse

    #parsing comand line argumetns
    parser = argparse.ArgumentParser(description='Convert BLAST xml file to FASTA file.')
    parser.add_argument('-i', '--input', nargs=1, help='Path to the input BLAST xml file. Default is stdin.', dest='input')
    parser.add_argument('-o', '--output', nargs=1,  help='Path to the output FASTA file. Default is stdout.', dest='output')
    args = parser.parse_args()

    #set input
    infile = sys.stdin
    if args.input != None:
        infilePath = args.input[0]
        try:
            infile = open(infilePath, 'r')
        except IOError:
            exit("Input file {0} cann't be opened!\n".format(infilePath))

    #set output
    openedOutFile = sys.stdout
    if args.output != None:
        outfilePath = args.output[0]
        try:
            openedOutFile = open(outfilePath, 'w')
        except IOError:
            exit("Output file {0} cann't be opened!\n".format(outfilePath))

    #read xml file to string
    try:
        tree = infile.read()
    except IOError:
        sys.stderr.write("ERRRO while reading from {0}\n".format(infile.name))
        exit(1)

    #convert xml to fasta and write to the output file
    try:
        fastaSeqs = convertXMLtoFASTA(tree)
    except ElementTree.ParseError:
        sys.stderr.write("ERROR in xml tree. Parser cann't parse it!\n")
        fastaSeqs = None

    infile.close()

    if fastaSeqs != None:
        openedOutFile.write(fastaSeqs)
    openedOutFile.close()
    if fastaSeqs == None:
        exit(1)
    exit(0)

